import numpy as np
import open3d as o3d
import os
import time
import sys
import concurrent.futures
from collections import defaultdict
from scipy.spatial import cKDTree
import multiprocessing

# ------ 参数 ---------
pcd_paths = [
    "../maps/partmap1.pcd",
    "../maps/partmap2.pcd",
]

output_path = "../maps/merged_map.pcd"
output_voxel_size = 0.02  # 体素大小
min_occurrences = 2  # 最小出现次数（用于动态物体去除）
if_visualize_when_finish = False  # 是否可视化点云
use_gpu = True  # 是否使用GPU加速
# ---------------------


def load_point_clouds(paths):
    """加载点云并打印基本信息"""
    pcds = []
    print("\n加载点云:")
    for i, path in enumerate(paths):
        if not os.path.exists(path):
            print(f"错误: 文件不存在 - {path}")
            sys.exit(1)

        pcd = o3d.io.read_point_cloud(path)
        if not pcd.has_points():
            print(f"错误: 点云为空 - {path}")
            sys.exit(1)

        print(f"  PCD {i+1}: {path}, 点数: {len(pcd.points)}")
        pcds.append(pcd)
    return pcds


def preprocess_point_cloud(pcd, voxel_size):
    """点云预处理：降采样和法线估计"""
    # 先粗略下采样减少计算量
    rough_down = pcd.voxel_down_sample(voxel_size * 4)

    # 检查点云是否为空
    if len(rough_down.points) == 0:
        print("警告: 降采样后点云为空，使用原始点云")
        pcd_down = pcd
    else:
        pcd_down = rough_down

    # 进一步下采样用于配准
    pcd_down = pcd_down.voxel_down_sample(voxel_size)

    # 估计法线 - 用于点对面ICP
    radius = voxel_size * 2
    pcd_down.estimate_normals(
        o3d.geometry.KDTreeSearchParamHybrid(radius=radius, max_nn=30)
    )

    return pcd_down


def pairwise_registration(source, target, voxel_size):
    """点云对配准"""
    distance_threshold = voxel_size * 1.5

    # 使用点对面ICP
    result = o3d.pipelines.registration.registration_icp(
        source,
        target,
        distance_threshold,
        np.identity(4),
        o3d.pipelines.registration.TransformationEstimationPointToPlane(),
        o3d.pipelines.registration.ICPConvergenceCriteria(
            max_iteration=50
        ),  # 减少迭代次数
    )
    return result.transformation


def multiway_registration(pcds, voxel_size):
    """多视角点云配准"""
    print("开始两两配准...")
    n = len(pcds)

    # 创建位姿图
    pose_graph = o3d.pipelines.registration.PoseGraph()

    # 添加节点
    for i in range(n):
        pose_graph.nodes.append(
            o3d.pipelines.registration.PoseGraphNode(np.identity(4))
        )

    # 只配准相邻的点云对，避免全连接
    for i in range(n - 1):
        j = i + 1
        print(f"  配准点云 {i+1} 和 {j+1}...")
        transformation = pairwise_registration(pcds[i], pcds[j], voxel_size)
        pose_graph.edges.append(
            o3d.pipelines.registration.PoseGraphEdge(
                i, j, transformation, uncertain=False, confidence=1.0
            )
        )

    print("进行全局优化...")
    # 全局优化参数
    option = o3d.pipelines.registration.GlobalOptimizationOption(
        max_correspondence_distance=voxel_size * 1.5,
        edge_prune_threshold=0.25,
        reference_node=0,
    )

    # 执行全局优化
    o3d.pipelines.registration.global_optimization(
        pose_graph,
        o3d.pipelines.registration.GlobalOptimizationLevenbergMarquardt(),
        o3d.pipelines.registration.GlobalOptimizationConvergenceCriteria(),
        option,
    )

    # 提取优化后的位姿
    optimized_poses = []
    for node in pose_graph.nodes:
        optimized_poses.append(node.pose)

    return optimized_poses


def compute_bounds(pcds):
    """计算所有点云的边界"""
    min_bound = np.array([np.inf, np.inf, np.inf])
    max_bound = np.array([-np.inf, -np.inf, -np.inf])

    for pcd in pcds:
        points = np.asarray(pcd.points)
        if len(points) == 0:
            continue

        min_bound = np.minimum(min_bound, np.min(points, axis=0))
        max_bound = np.maximum(max_bound, np.max(points, axis=0))

    return min_bound, max_bound


def voxel_index(point, min_bound, voxel_size):
    """计算点的体素索引"""
    return tuple(((point - min_bound) // voxel_size).astype(int))


def process_point_cloud(args):
    """处理单个点云的函数（用于并行）"""
    cloud_idx, pcd, min_bound, voxel_size = args
    points = np.asarray(pcd.points)
    cloud_voxels = defaultdict(set)

    for point in points:
        idx = voxel_index(point, min_bound, voxel_size)
        cloud_voxels[idx].add(tuple(point))

    return cloud_idx, cloud_voxels


def remove_dynamic_objects_optimized(registered_pcds, voxel_size, min_occurrences):
    """优化后的去除动态物体函数 - 使用并行处理和高效数据结构"""
    print("优化版去除动态物体...")
    start_time = time.time()

    # 1. 计算整体边界
    min_bound, max_bound = compute_bounds(registered_pcds)
    print(f"  场景边界: min={min_bound}, max={max_bound}")

    # 2. 并行处理每个点云
    print("  并行处理点云...")
    args_list = [
        (i, pcd, min_bound, voxel_size) for i, pcd in enumerate(registered_pcds)
    ]

    # 使用多进程并行处理
    with concurrent.futures.ProcessPoolExecutor() as executor:
        results = list(executor.map(process_point_cloud, args_list))

    # 3. 合并体素信息
    print("  合并体素信息...")
    global_voxels = defaultdict(set)
    cloud_voxel_sets = {}

    for cloud_idx, cloud_voxels in results:
        cloud_voxel_sets[cloud_idx] = cloud_voxels
        for voxel_idx, points in cloud_voxels.items():
            global_voxels[voxel_idx].add(cloud_idx)

    # 4. 收集静态点
    print("  收集静态点...")
    static_points = []

    # 筛选出现次数足够的体素
    for voxel_idx, cloud_ids in global_voxels.items():
        if len(cloud_ids) >= min_occurrences:
            # 从所有点云中收集该体素内的点
            for cloud_idx in cloud_ids:
                if voxel_idx in cloud_voxel_sets[cloud_idx]:
                    static_points.extend(cloud_voxel_sets[cloud_idx][voxel_idx])

    # 5. 创建最终点云
    static_pcd = o3d.geometry.PointCloud()
    static_pcd.points = o3d.utility.Vector3dVector(np.array(static_points))

    print(f"  优化版去除动态物体完成! 耗时: {time.time() - start_time:.2f}秒")
    return static_pcd


# 全局函数，解决Windows平台pickle问题
def process_point_chunk(chunk_points, ref_points, voxel_size, min_occurrences):
    """处理点云分块，用于并行处理"""
    if len(ref_points) == 0:
        # 如果参考点为空，则返回空列表
        return []

    static_chunk = []

    # 在子进程内构建KDTree
    local_kdtree = cKDTree(ref_points)

    # 批量查询所有点的邻居数量
    neighbor_counts = local_kdtree.query_ball_point(
        chunk_points, r=voxel_size * 6, return_length=True
    )

    # 筛选满足最小出现次数的点
    for i, count in enumerate(neighbor_counts):
        if count >= min_occurrences:
            static_chunk.append(chunk_points[i])

    return static_chunk


def remove_dynamic_objects_fast_optimized(registered_pcds, voxel_size, min_occurrences):
    """优化版快速去除动态物体方法 - 使用KDTree和并行处理"""
    print("优化版快速去除动态物体...")
    start_time = time.time()

    # 1. 创建参考点云（所有点云的组合）并下采样
    all_points = np.vstack([np.asarray(pcd.points) for pcd in registered_pcds])
    reference_pcd = o3d.geometry.PointCloud()
    reference_pcd.points = o3d.utility.Vector3dVector(all_points)

    # 更激进的下采样以减少计算量
    downsampled_ref = reference_pcd.voxel_down_sample(voxel_size * 8)
    ref_points = np.asarray(downsampled_ref.points)

    # 检查参考点云是否为空
    if len(ref_points) == 0:
        print("警告: 参考点云下采样后为空，使用原始点云")
        ref_points = all_points

    print(f"  参考点云下采样: {len(all_points)} -> {len(ref_points)} 点")

    # 2. 并行处理点云块
    print("  并行处理点云块...")
    static_points = []

    # 使用多进程并行处理
    with concurrent.futures.ProcessPoolExecutor(
        max_workers=min(8, multiprocessing.cpu_count())
    ) as executor:
        futures = []
        for pcd in registered_pcds:
            points = np.asarray(pcd.points)

            # 跳过空点云
            if len(points) == 0:
                continue

            # 将点云分割成块以便并行处理
            chunk_size = max(50000, len(points) // (multiprocessing.cpu_count() * 2))
            chunks = [
                points[i : i + chunk_size] for i in range(0, len(points), chunk_size)
            ]

            # 提交任务
            for chunk in chunks:
                futures.append(
                    executor.submit(
                        process_point_chunk,  # 使用全局函数
                        chunk,
                        ref_points,
                        voxel_size,
                        min_occurrences,
                    )
                )

        # 收集结果并显示进度
        completed = 0
        total = len(futures)
        for future in concurrent.futures.as_completed(futures):
            try:
                result = future.result()
                static_points.extend(result)
            except Exception as e:
                print(f"  处理块时出错: {e}")

            completed += 1
            if completed % 10 == 0 or completed == total:
                print(
                    f"  处理进度: {completed}/{total} 块 ({completed/total*100:.1f}%)"
                )

    # 3. 创建最终点云
    static_pcd = o3d.geometry.PointCloud()
    if static_points:
        # 使用vstack更高效地组合点
        static_pcd.points = o3d.utility.Vector3dVector(np.vstack(static_points))
    else:
        print("警告: 未找到静态点，创建空点云")
        static_pcd.points = o3d.utility.Vector3dVector(np.empty((0, 3)))

    print(
        f"  优化版完成! 耗时: {time.time() - start_time:.2f}秒, 静态点: {len(static_pcd.points)}"
    )
    return static_pcd


def main():
    print("开始点云融合处理...")
    start_time = time.time()

    # 1. 加载点云
    pcds = load_point_clouds(pcd_paths)

    # 2. 预处理 - 使用较大的体素减少数据量
    print("\n预处理点云...")
    voxel_size = output_voxel_size * 4  # 配准使用较大的体素
    pcds_down = []
    for i, pcd in enumerate(pcds):
        print(f"  预处理点云 {i+1}...")
        pcd_down = preprocess_point_cloud(pcd, voxel_size)
        pcds_down.append(pcd_down)
        print(f"    降采样后点数: {len(pcd_down.points)}")

    # 3. 多视角配准
    print("\n开始多视角配准...")
    poses = multiway_registration(pcds_down, voxel_size)

    # 4. 应用变换到原始点云
    print("\n应用变换到原始点云...")
    registered_pcds = []
    for i, pcd in enumerate(pcds):
        print(f"  应用变换到点云 {i+1}...")
        transformed_pcd = pcd.transform(poses[i])
        registered_pcds.append(transformed_pcd)

    # 5. 去除动态物体并融合
    print("\n去除动态物体并融合点云...")

    # 根据点云大小选择合适的算法
    total_points = sum(len(pcd.points) for pcd in registered_pcds)
    print(f"  总点数: {total_points}")

    # 使用优化后的快速方法
    static_pcd = remove_dynamic_objects_fast_optimized(
        registered_pcds, output_voxel_size, min_occurrences
    )

    print(f"  静态点数量: {len(static_pcd.points)}")

    # 6. 体素下采样
    print("执行体素下采样...")
    if len(static_pcd.points) > 0:
        final_pcd = static_pcd.voxel_down_sample(output_voxel_size)
    else:
        print("警告: 静态点云为空，使用原始点云")
        final_pcd = static_pcd

    # 7. 后处理：去噪
    print("去除统计离群点...")
    if len(final_pcd.points) > 0:
        final_pcd, _ = final_pcd.remove_statistical_outlier(
            nb_neighbors=20, std_ratio=2.0
        )

    # 8. 保存结果
    print(f"\n保存结果到 {output_path}")
    o3d.io.write_point_cloud(output_path, final_pcd)

    # 9. 可视化
    if if_visualize_when_finish:
        print("可视化结果...")
        o3d.visualization.draw_geometries([final_pcd])

    # 打印统计信息
    print("\n处理完成!")
    print(f"  输入点云数量: {len(pcds)}")
    print(f"  输入总点数: {sum(len(pcd.points) for pcd in pcds)}")
    print(f"  输出点数: {len(final_pcd.points)}")
    print(f"  处理时间: {time.time() - start_time:.2f}秒")


if __name__ == "__main__":
    # 设置多进程启动方法
    if sys.platform.startswith("win"):
        multiprocessing.set_start_method("spawn", force=True)
    else:
        multiprocessing.set_start_method("fork", force=True)
    main()
